import { create } from 'zustand'
import { sessionService, SessionResponse, SessionFilters } from '../services/sessionService'
import { RoleType } from '../components/ui/RoleSelector'

export interface Message {
  id: string
  role: 'user' | 'assistant' | 'system'
  content: string
  timestamp: Date
  sequence_number: number
  branch_id?: string
  metadata?: {
    model?: string
    temperature?: number
    searchResults?: any[]
  }
}

export interface Session {
  id: string
  title: string
  role_id?: string
  model_config?: {
    model: string
    temperature: number
    max_tokens: number
  }
  settings?: {
    searchEnabled: boolean
    contextThreshold: number
    autoSummary: boolean
  }
  is_favorite?: boolean
  is_archived?: boolean
  message_count?: number
  last_message_at?: string
  created_at: Date
  updated_at: Date
}

export interface Role {
  id: string
  name: string
  description?: string
  system_prompt: string
  model_preferences?: {
    preferredModel: string
    temperature: number
    max_tokens: number
    top_p: number
  }
  is_system: boolean
}

interface ChatState {
  // Sessions
  sessions: Session[]
  currentSession: Session | null
  
  // 分页和筛选
  pagination: {
    page: number
    pageSize: number
    total: number
    totalPages: number
  }
  filters: {
    query: string
    sort: 'lastActiveAt' | 'title' | 'created_at'
    order: 'asc' | 'desc'
    filter: 'favorite' | 'archived' | ''
  }
  
  // Messages
  messages: Message[]
  isStreaming: boolean
  
  // Roles
  roles: Role[]
  currentRole: RoleType
  
  // Loading状态
  isLoading: boolean
  
  // UI State
  sidebarOpen: boolean
  
  // WebSocket
  ws: WebSocket | null
  connectionStatus: 'connecting' | 'connected' | 'disconnected'
  
  // Session Actions
  loadSessions: () => Promise<void>
  searchSessions: (query: string) => Promise<void>
  loadMoreSessions: () => Promise<void>
  createSession: (init?: { title?: string; roleId?: string; model?: string }) => Promise<Session>
  renameSession: (sessionId: string, title: string) => Promise<void>
  deleteSession: (sessionId: string) => Promise<void>
  restoreSession: (sessionId: string) => Promise<void>
  toggleFavoriteSession: (sessionId: string) => Promise<void>
  toggleArchiveSession: (sessionId: string) => Promise<void>
  setCurrentSession: (session: Session | null) => void
  updateSessionSettings: (sessionId: string, settings: {
    roleId?: string
    model?: string
    searchEnabled?: boolean
  }) => Promise<void>
  
  // 筛选和排序
  updateFilters: (filters: Partial<{
    query: string
    sort: 'lastActiveAt' | 'title' | 'created_at'
    order: 'asc' | 'desc'
    filter: 'favorite' | 'archived' | ''
  }>) => Promise<void>
  
  // Legacy Actions (保持兼容)
  setSessions: (sessions: Session[]) => void
  
  setMessages: (messages: Message[]) => void
  addMessage: (message: Message) => void
  updateMessage: (messageId: string, content: string) => void
  
  setRoles: (roles: Role[]) => void
  setCurrentRole: (role: Role | null) => void
  
  setSidebarOpen: (open: boolean) => void
  setIsStreaming: (streaming: boolean) => void
  
  connectWebSocket: (sessionId: string) => void
  disconnectWebSocket: () => void
  sendMessage: (content: string) => void
}

export const useChatStore = create<ChatState>((set, get) => ({
  // Initial state
  sessions: [],
  currentSession: null,
  pagination: {
    page: 1,
    pageSize: 20,
    total: 0,
    totalPages: 0
  },
  filters: {
    query: '',
    sort: 'lastActiveAt',
    order: 'desc',
    filter: ''
  },
  messages: [],
  isStreaming: false,
  roles: [],
  currentRole: 'emma',
  isLoading: false,
  sidebarOpen: true,
  ws: null,
  connectionStatus: 'disconnected',

  // Session actions
  loadSessions: async () => {
    set({ isLoading: true })
    try {
      const { filters, pagination } = get()
      const response = await sessionService.getSessions({
        ...filters,
        page: pagination.page,
        pageSize: pagination.pageSize
      })
      
      const sessions = response.sessions.map(s => ({
        ...s,
        created_at: new Date(s.created_at),
        updated_at: new Date(s.updated_at)
      }))
      
      set({
        sessions,
        pagination: response.pagination,
        isLoading: false
      })
    } catch (error) {
      console.error('Failed to load sessions:', error)
      set({ isLoading: false })
    }
  },

  searchSessions: async (query: string) => {
    set(state => ({
      filters: { ...state.filters, query },
      pagination: { ...state.pagination, page: 1 }
    }))
    await get().loadSessions()
  },

  loadMoreSessions: async () => {
    const { pagination } = get()
    if (pagination.page >= pagination.totalPages) return
    
    set(state => ({
      pagination: { ...state.pagination, page: state.pagination.page + 1 }
    }))
    
    try {
      const { filters, pagination: newPagination } = get()
      const response = await sessionService.getSessions({
        ...filters,
        page: newPagination.page,
        pageSize: newPagination.pageSize
      })
      
      const newSessions = response.sessions.map(s => ({
        ...s,
        created_at: new Date(s.created_at),
        updated_at: new Date(s.updated_at)
      }))
      
      set(state => ({
        sessions: [...state.sessions, ...newSessions],
        pagination: response.pagination
      }))
    } catch (error) {
      console.error('Failed to load more sessions:', error)
    }
  },

  createSession: async (init = {}) => {
    try {
      const { currentRole } = get()
      const response = await sessionService.createSession({
        title: init.title || '新对话',
        role_id: init.roleId || currentRole,
        model: init.model || 'doubao-seed-1-6-250615',
        searchEnabled: true
      })
      
      const newSession: Session = {
        ...response,
        created_at: new Date(response.created_at),
        updated_at: new Date(response.updated_at)
      }
      
      set(state => ({
        sessions: [newSession, ...state.sessions],
        currentSession: newSession,
        messages: []
      }))
      
      return newSession
    } catch (error) {
      console.error('Failed to create session:', error)
      throw error
    }
  },

  renameSession: async (sessionId: string, title: string) => {
    try {
      const response = await sessionService.updateSession(sessionId, { title })
      
      set(state => ({
        sessions: state.sessions.map(s => 
          s.id === sessionId 
            ? { ...s, title, updated_at: new Date(response.updated_at) }
            : s
        ),
        currentSession: state.currentSession?.id === sessionId 
          ? { ...state.currentSession, title }
          : state.currentSession
      }))
    } catch (error) {
      console.error('Failed to rename session:', error)
      throw error
    }
  },

  deleteSession: async (sessionId: string) => {
    try {
      await sessionService.deleteSession(sessionId)
      
      set(state => ({
        sessions: state.sessions.filter(s => s.id !== sessionId),
        currentSession: state.currentSession?.id === sessionId ? null : state.currentSession
      }))
    } catch (error) {
      console.error('Failed to delete session:', error)
      throw error
    }
  },

  restoreSession: async (sessionId: string) => {
    try {
      const response = await sessionService.restoreSession(sessionId)
      // 重新加载会话列表
      await get().loadSessions()
    } catch (error) {
      console.error('Failed to restore session:', error)
      throw error
    }
  },

  toggleFavoriteSession: async (sessionId: string) => {
    try {
      const response = await sessionService.toggleFavorite(sessionId)
      
      set(state => ({
        sessions: state.sessions.map(s => 
          s.id === sessionId 
            ? { ...s, is_favorite: response.is_favorite }
            : s
        )
      }))
    } catch (error) {
      console.error('Failed to toggle favorite:', error)
      throw error
    }
  },

  toggleArchiveSession: async (sessionId: string) => {
    try {
      const response = await sessionService.toggleArchive(sessionId)
      
      set(state => ({
        sessions: state.sessions.map(s => 
          s.id === sessionId 
            ? { ...s, is_archived: response.is_archived }
            : s
        )
      }))
    } catch (error) {
      console.error('Failed to toggle archive:', error)
      throw error
    }
  },

  setCurrentSession: (session: Session | null) => {
    set({ currentSession: session, messages: [] })
    // TODO: Load messages for this session
  },

  updateSessionSettings: async (sessionId: string, settings: {
    roleId?: string
    model?: string
    searchEnabled?: boolean
  }) => {
    try {
      const updateData: any = {}
      if (settings.roleId) updateData.role_id = settings.roleId
      if (settings.model) {
        updateData.model_config = {
          model: settings.model,
          temperature: 0.7,
          max_tokens: 1500
        }
      }
      if (settings.searchEnabled !== undefined) {
        updateData.settings = { searchEnabled: settings.searchEnabled }
      }

      const response = await sessionService.updateSession(sessionId, updateData)
      
      set(state => ({
        sessions: state.sessions.map(s => 
          s.id === sessionId ? { ...s, ...updateData, updated_at: new Date(response.updated_at) } : s
        ),
        currentSession: state.currentSession?.id === sessionId 
          ? { ...state.currentSession, ...updateData }
          : state.currentSession
      }))
    } catch (error) {
      console.error('Failed to update session settings:', error)
      throw error
    }
  },

  updateFilters: async (newFilters) => {
    set(state => ({
      filters: { ...state.filters, ...newFilters },
      pagination: { ...state.pagination, page: 1 } // 重置分页
    }))
    await get().loadSessions()
  },

  // Legacy Actions (保持兼容)
  setSessions: (sessions) => set({ sessions }),

  // Message actions
  setMessages: (messages) => set({ messages }),
  
  addMessage: (message) => {
    set(state => ({ 
      messages: [...state.messages, message] 
    }))
  },
  
  updateMessage: (messageId, content) => {
    set(state => ({
      messages: state.messages.map(msg => 
        msg.id === messageId ? { ...msg, content } : msg
      )
    }))
  },

  // Role actions
  setRoles: (roles) => set({ roles }),
  setCurrentRole: (role) => set({ currentRole: role }),

  // UI actions
  setSidebarOpen: (open) => set({ sidebarOpen: open }),
  setIsStreaming: (streaming) => set({ isStreaming: streaming }),

  // WebSocket actions
  connectWebSocket: (sessionId: string) => {
    const ws = new WebSocket(`ws://localhost:8004/api/v1/ws/${sessionId}`)
    
    ws.onopen = () => {
      set({ ws, connectionStatus: 'connected' })
      console.log('WebSocket connected')
    }
    
    ws.onmessage = (event) => {
      const data = JSON.parse(event.data)
      
      if (data.type === 'message') {
        get().addMessage(data.message)
      } else if (data.type === 'stream_chunk') {
        // Handle streaming message updates
        const lastMessage = get().messages[get().messages.length - 1]
        if (lastMessage && lastMessage.role === 'assistant') {
          get().updateMessage(lastMessage.id, lastMessage.content + data.content)
        }
      } else if (data.type === 'stream_start') {
        get().setIsStreaming(true)
        // Add empty assistant message for streaming
        const streamMessage: Message = {
          id: data.message_id,
          role: 'assistant',
          content: '',
          timestamp: new Date(),
          sequence_number: data.sequence_number
        }
        get().addMessage(streamMessage)
      } else if (data.type === 'stream_end') {
        get().setIsStreaming(false)
      }
    }
    
    ws.onclose = () => {
      set({ ws: null, connectionStatus: 'disconnected' })
      console.log('WebSocket disconnected')
    }
    
    ws.onerror = (error) => {
      console.error('WebSocket error:', error)
      set({ connectionStatus: 'disconnected' })
    }
    
    set({ connectionStatus: 'connecting' })
  },
  
  disconnectWebSocket: () => {
    const { ws } = get()
    if (ws) {
      ws.close()
      set({ ws: null, connectionStatus: 'disconnected' })
    }
  },
  
  sendMessage: (content: string) => {
    const { ws, currentSession } = get()
    
    if (!ws || !currentSession) {
      console.error('WebSocket not connected or no active session')
      return
    }
    
    // Add user message immediately
    const userMessage: Message = {
      id: crypto.randomUUID(),
      role: 'user',
      content,
      timestamp: new Date(),
      sequence_number: get().messages.length + 1
    }
    
    get().addMessage(userMessage)
    
    // Send to WebSocket
    ws.send(JSON.stringify({
      type: 'user_message',
      content,
      session_id: currentSession.id,
      role_id: get().currentRole?.id
    }))
  }
}))